package com.by.basicgateway.balancer;

import cn.hutool.core.date.StopWatch;
import com.by.toolkit.consts.ByContextConst;
import com.google.common.collect.Maps;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.ObjectProvider;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.loadbalancer.Request;
import org.springframework.cloud.client.loadbalancer.RequestDataContext;
import org.springframework.cloud.client.loadbalancer.Response;
import org.springframework.cloud.loadbalancer.core.ServiceInstanceListSupplier;
import reactor.core.publisher.Mono;

import java.util.List;
import java.util.Map;
import java.util.Random;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ScheduledThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.Collectors;

import static com.by.toolkit.utils.ExceptionUtil.getStackTrace;

/**
 * 重写
 *
 * @author Tian Wei
 * @version 1.0
 * @date 2024/2/18 7:15 PM
 */
@Slf4j
public class ByRoundRobinLoadBalancer extends org.springframework.cloud.loadbalancer.core.RoundRobinLoadBalancer {

    final AtomicInteger position;

    final String serviceId;

    private static final Map<String, Long> requestIdsMap = Maps.newConcurrentMap();

    private static final ScheduledExecutorService SCHEDULED_EXECUTOR_SERVICE = new ScheduledThreadPoolExecutor(1);

    ObjectProvider<ServiceInstanceListSupplier> serviceInstanceListSupplierProvider;

    static {
        SCHEDULED_EXECUTOR_SERVICE.scheduleWithFixedDelay(() -> {
            try {
                StopWatch stopWatch = new StopWatch();
                stopWatch.start("开始执行清除任务");
                List<String> waitDeletedKeys = requestIdsMap.entrySet().stream()
                        .filter(item -> item.getValue() < System.currentTimeMillis())
                        .map(Map.Entry::getKey)
                        .collect(Collectors.toList());
                for (String waitDeletedKey : waitDeletedKeys) {
                    requestIdsMap.remove(waitDeletedKey);
                }
                stopWatch.stop();
                log.info("[ByRoundRobinLoadBalancer][afterPropertiesSet] 清除任务执行完成,执行时间为:{},清除数量:{}", stopWatch.getTotalTimeMillis(), waitDeletedKeys.size());

            } catch (Exception e) {
                log.info("[ByRoundRobinLoadBalancer][static initializer] 清除任务执行异常:{}", getStackTrace(e));
            }
        }, 20, 60, TimeUnit.SECONDS);
    }

    public ByRoundRobinLoadBalancer(ObjectProvider<ServiceInstanceListSupplier> serviceInstanceListSupplierProvider,
                                    String serviceId) {
        this(serviceInstanceListSupplierProvider, serviceId, new Random().nextInt(1000));
    }

    public ByRoundRobinLoadBalancer(ObjectProvider<ServiceInstanceListSupplier> serviceInstanceListSupplierProvider,
                                    String serviceId, int seedPosition) {
        super(serviceInstanceListSupplierProvider, serviceId, seedPosition);
        this.serviceId = serviceId;
        this.serviceInstanceListSupplierProvider = serviceInstanceListSupplierProvider;
        this.position = new AtomicInteger(seedPosition);
    }

    @SuppressWarnings("rawtypes")
    @Override
    public Mono<Response<ServiceInstance>> choose(Request request) {

        //强制转换,获取requestId
        RequestDataContext context = (RequestDataContext) request.getContext();
        String requestId = context.getClientRequest().getHeaders().getFirst(ByContextConst.REQUEST_ID);
        if (requestIdsMap.containsKey(requestId)) {
            return Mono.empty();
        } else {
            requestIdsMap.put(requestId, System.currentTimeMillis() + 30 * 1000L);
        }

        return super.choose(request);
    }

}
